from app.tasks import celery
from app.services.document_service import DocumentService
from app.utils.errors import ServiceError
from app.models.task import Task, TaskStatus, TaskType
from app import logger
from flask import current_app
import traceback
from celery import Task
from app.services.notification_service import NotificationService
from celery.utils.log import get_task_logger

celery_logger = get_task_logger(__name__)

class BaseTask(Task):
    """基础任务类"""
    abstract = True
    
    def on_failure(self, exc, task_id, args, kwargs, einfo):
        """任务失败处理"""
        NotificationService.notify('task_failure', {
            'task_id': task_id,
            'task_name': self.name,
            'error': str(exc),
            'args': args,
            'kwargs': kwargs
        })
        super().on_failure(exc, task_id, args, kwargs, einfo)
    
    def on_success(self, retval, task_id, args, kwargs):
        """任务成功处理"""
        NotificationService.notify('task_success', {
            'task_id': task_id,
            'task_name': self.name,
            'result': retval,
            'args': args,
            'kwargs': kwargs
        })
        super().on_success(retval, task_id, args, kwargs)

@celery.task(
    bind=True,
    base=BaseTask,
    max_retries=3,
    default_retry_delay=60,  # 60秒后重试
    priority=5  # 默认优先级
)
def process_document(self, file_path: str, knowledge_base_id: int):
    """处理文档的异步任务"""
    task = Task.create_task(self.request.id, TaskType.DOCUMENT_PROCESS)
    
    try:
        task.update_status(TaskStatus.RUNNING, progress=0)
        
        logger.info(f"Starting document processing: {file_path}")
        doc = DocumentService.process_document(file_path, knowledge_base_id)
        
        result = {'document_id': doc.id}
        task.update_status(TaskStatus.SUCCESS, progress=100, result=result)
        
        return result
        
    except Exception as e:
        error_msg = f"Document processing failed: {str(e)}"
        logger.error(f"{error_msg}\n{traceback.format_exc()}")
        task.update_status(TaskStatus.FAILURE, error=error_msg)
        
        # 重试任务
        try:
            raise self.retry(exc=e)
        except self.MaxRetriesExceededError:
            raise ServiceError(error_msg)

@celery.task(
    bind=True,
    base=BaseTask,
    max_retries=3,
    default_retry_delay=30,  # 30秒后重试
    priority=7  # 较高优先级
)
def generate_embeddings(self, document_id: int):
    """生成文档向量的异步任务"""
    task = Task.create_task(self.request.id, TaskType.EMBEDDING_GENERATE)
    
    try:
        task.update_status(TaskStatus.RUNNING, progress=0)
        
        logger.info(f"Generating embeddings for document: {document_id}")
        DocumentService.generate_embeddings(document_id)
        
        result = {'document_id': document_id}
        task.update_status(TaskStatus.SUCCESS, progress=100, result=result)
        
        return result
        
    except Exception as e:
        error_msg = f"Embedding generation failed: {str(e)}"
        logger.error(f"{error_msg}\n{traceback.format_exc()}")
        task.update_status(TaskStatus.FAILURE, error=error_msg)
        raise ServiceError(error_msg) 